#include <stdio.h>
#include <netdb.h>
#include <stdlib.h>
#include <string.h>
#include <unistd.h>
#include <rdma/rdma_cma.h>
#include <infiniband/verbs.h>

const int BUFFER_SIZE = 1024;

struct context {
    struct ibv_context *ctx;
    struct ibv_pd *pd;
    struct ibv_cq *cq;
    struct ibv_comp_channel *comp_channel;
    pthread_t cq_poller_thread;
};

static struct context *s_ctx = NULL;

struct connection {
    struct ibv_qp *qp;
    struct ibv_mr *recv_mr;
    struct ibv_mr *send_mr;
    char *recv_region;
    char *send_region;

    struct rdma_cm_id *id;
    int num_completions;
};

void *
poll_cq(void *ctx)
{
    struct ibv_cq *cq;
    struct ibv_wc wc;
    while (1) {
        ibv_get_cq_event(s_ctx->comp_channel, &cq, &ctx);
        ibv_ack_cq_events(cq, 1);
        ibv_req_notify_cq(cq, 0);

        while (ibv_poll_cq(cq, 1, &wc)) {
            struct connection *conn = (struct connection *)(uintptr_t)(wc.wr_id);
            if (wc.status != IBV_WC_SUCCESS) {
                printf("on_completion: status is not IBV_WC_SUCCESS.");
            }

            if (wc.opcode & IBV_WC_RECV) {
                printf("received message: %s\n", conn->recv_region);
            } else if (wc.opcode == IBV_WC_SEND) {
                printf("send completed successfully.\n");
            }

            /*这个在客户端调用了，服务端就不用了*/
            /*if (++conn->num_completions == 2){
                rdma_disconnect(conn->id);
                printf("server disconnect\n");
            }*/

        }
    }
    return NULL;
}

int connect_request(struct rdma_cm_id *id)
{
    struct ibv_qp_init_attr qp_attr;
    struct rdma_conn_param cm_params;
    struct ibv_recv_wr wr, *bad_wr = NULL;
    struct ibv_sge sge;
    struct connection *conn;

    if (s_ctx) {
        if (s_ctx->ctx != id->verbs)
          printf("cannot handle events in more than one context.");
        return -1;
    }

    s_ctx = (struct context *)malloc(sizeof(struct context));
    if(!s_ctx){
        printf("malloc error\n");
        return -1;
    }

    s_ctx->ctx = id->verbs;
    s_ctx->pd = ibv_alloc_pd(s_ctx->ctx);
    if(!s_ctx->pd){
        printf("ibv_alloc_pd error\n");
    }

    s_ctx->comp_channel = ibv_create_comp_channel(s_ctx->ctx);
    if(!s_ctx->comp_channel){
        printf("ibv_create_comp_channel error\n");
    }
    s_ctx->cq = ibv_create_cq(s_ctx->ctx, 10, NULL, s_ctx->comp_channel, 0);
    if(!s_ctx->cq){
        printf("ibv_create_cq error\n");
    }

    if(ibv_req_notify_cq(s_ctx->cq, 0)){
        printf("ibv_req_notify_cq error\n");
    }

    pthread_create(&s_ctx->cq_poller_thread, NULL, poll_cq, NULL);

    memset(&qp_attr, 0, sizeof(qp_attr));
    qp_attr.send_cq = s_ctx->cq;
    qp_attr.recv_cq = s_ctx->cq;
    qp_attr.qp_type = IBV_QPT_RC;
    qp_attr.cap.max_send_wr = 10;
    qp_attr.cap.max_recv_wr = 10;
    qp_attr.cap.max_send_sge = 1;
    qp_attr.cap.max_recv_sge = 1;

    if(rdma_create_qp(id, s_ctx->pd, &qp_attr)) {
        printf("rdma_create_qp error\n");
    }

    id->context = conn = (struct connection *)malloc(sizeof(struct connection));

    conn->qp = id->qp;
    conn->id = id;
    conn->num_completions = 0;

    conn->send_region = malloc(BUFFER_SIZE);
    conn->recv_region = malloc(BUFFER_SIZE);
    conn->send_mr = ibv_reg_mr(s_ctx->pd, conn->send_region, BUFFER_SIZE, 0);
    if(!conn->send_mr) {
        printf("ibv_reg_mr error\n");
    }
    conn->recv_mr = ibv_reg_mr(s_ctx->pd, conn->recv_region, BUFFER_SIZE, IBV_ACCESS_LOCAL_WRITE);
    if(!conn->recv_mr){
        printf("ibv_reg_mr error\n");
    }

    sge.addr = (uintptr_t)conn->recv_region;
    sge.length = BUFFER_SIZE;
    sge.lkey = conn->recv_mr->lkey;

    wr.wr_id = (uintptr_t)conn;
    wr.next = NULL;
    wr.sg_list = &sge;
    wr.num_sge = 1;
    if(ibv_post_recv(conn->qp, &wr, &bad_wr)) {
        printf("ibv_post_recv error\n");
    }

    memset(&cm_params, 0, sizeof(cm_params));
    /*走到这里，应该就会发起建立连接产生的第二个包，客户端收到这个包以后
    应该会产生一个RDMA_CM_EVENT_ESTABLISHED事件，然后回CM ReadyToUse这个包，
    服务段应该就会收到RDMA_CM_EVENT_ESTABLISHED事件？只是猜测*/
    if(rdma_accept(id, &cm_params)) {
        printf("rdma_accept error\n");
    }
    return 0;
}

int
connect_event(void *context)
{
    struct connection *conn = (struct connection *)context;
    struct ibv_send_wr wr, *bad_wr = NULL;
    struct ibv_sge sge;

    snprintf(conn->send_region, BUFFER_SIZE, "message from passive/server side with pid %d", getpid());
    memset(&wr, 0, sizeof(wr));

    sge.addr = (uintptr_t)conn->send_region;
    sge.length = BUFFER_SIZE;
    sge.lkey = conn->send_mr->lkey;

    wr.wr_id = (uintptr_t)conn;
    wr.opcode = IBV_WR_SEND /*IBV_WR_SEND_WITH_IMM*/;
    wr.sg_list = &sge;
    wr.num_sge = 1;
    wr.send_flags = IBV_SEND_SIGNALED;

    ibv_post_send(conn->qp, &wr, &bad_wr);
    return 0;
}

int
disconnect_event(struct rdma_cm_id *id)
{
    struct connection *conn = (struct connection *)id->context;

    rdma_destroy_qp(id);
    ibv_dereg_mr(conn->send_mr);
    ibv_dereg_mr(conn->recv_mr);
    free(conn->send_region);
    free(conn->recv_region);
    free(conn);
    rdma_destroy_id(id);
    return 0;
}

int
server_event_loop(struct rdma_cm_event *event)
{
    int ret = 0;
    /*客户端发起连接以后就会收到这个事件*/
    if(event->event == RDMA_CM_EVENT_CONNECT_REQUEST) {
        ret = connect_request(event->id);
        printf("RDMA_CM_EVENT_CONNECT_REQUEST\n");
    } else if(event->event == RDMA_CM_EVENT_ESTABLISHED) {
        /*服务端可以走到这里应该建立连接的过程已经ok了*/
        ret = connect_event(event->id->context);
        printf("RDMA_CM_EVENT_ESTABLISHED \n");
    } else if(event->event == RDMA_CM_EVENT_DISCONNECTED){
        ret = disconnect_event(event->id);
        printf("RDMA_CM_EVENT_DISCONNECTED \n");
    } else {
        printf("unknown event\n");
    }
    return ret;
}

int main(int argc, char **argv)
{
    struct sockaddr_in addr;
    struct rdma_cm_event *event = NULL;
    struct rdma_cm_id *listener = NULL;
    struct rdma_event_channel *ec = NULL;
    uint16_t port = 0, ret = 0;

    memset(&addr, 0, sizeof(addr));
    addr.sin_family = AF_INET;

    ec = rdma_create_event_channel();
    if(!ec) {
        printf("rdma_create_event_channel error\n");
        return -1;
    }

    ret = rdma_create_id(ec, &listener, NULL, RDMA_PS_TCP);
    if(ret) {
        printf("rdma_create_id error\n");
        return -1;
    }

    ret = rdma_bind_addr(listener, (struct sockaddr *)&addr);
    if(ret) {
        printf("rdma_bind_addr error\n");
        return -1;
    }

    ret = rdma_listen(listener, 10);
    if(ret) {
        printf("rdma_listen error\n");
        return -1;
    }

    port = ntohs(rdma_get_src_port(listener));
    printf("listening on port %d.\n", port);

    /*阻塞调用*/
    while (rdma_get_cm_event(ec, &event) == 0) {
        struct rdma_cm_event event_copy;
        memcpy(&event_copy, event, sizeof(*event));
        rdma_ack_cm_event(event);

        if (server_event_loop(&event_copy)) {
            break;
        }
    }

    rdma_destroy_id(listener);
    rdma_destroy_event_channel(ec);
    return 0;
}